import sqlite3

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.constants import START, END
from langgraph.graph import StateGraph
from langgraph.graph.state import CompiledStateGraph

from video_summarize_tools import nodes, edges
from video_summarize_tools.data_models import State, RuntimeContext

graph = StateGraph(State, context_schema=RuntimeContext)
graph.add_node("source_url_preprocess", nodes.source_url_preprocess)
graph.add_node("download_from_video_url", nodes.download_from_video_url)
graph.add_node("download_from_domain_url", nodes.download_from_domain_url)
graph.add_node("video_file_post_process", nodes.video_file_post_process)
graph.add_node("audio_transcript", nodes.audio_transcript)
graph.add_node("extract_video_keyframe", nodes.extract_video_keyframe)
graph.add_node("video_content_understanding", nodes.video_content_understanding)
graph.add_node("merge_result", nodes.merge_result)
graph.add_node("write_markdown_result_file", nodes.write_markdown_result_file)

graph.add_edge(START, "source_url_preprocess")
graph.add_conditional_edges("source_url_preprocess", edges.source_url_conditional_edge)
graph.add_edge("download_from_video_url", "video_file_post_process")
graph.add_edge("download_from_domain_url", "video_file_post_process")
graph.add_conditional_edges("video_file_post_process", edges.analysis_boot_edge)
graph.add_edge("extract_video_keyframe", "video_content_understanding")
graph.add_edge(["video_content_understanding", "audio_transcript"], "merge_result")
graph.add_edge("merge_result", "write_markdown_result_file")
graph.add_edge("write_markdown_result_file", END)
# 这里写了两个sub_task的edge，在merge_result这里会等待吗，还是继续执行

conn = sqlite3.connect("checkpoints.sqlite", check_same_thread=False)
compiled_graph = graph.compile(checkpointer=SqliteSaver(conn))

def get_graph() -> CompiledStateGraph:
    """
    获取预编译的graph
    :return:
    """
    return compiled_graph